"""
Tree可视化模块
提供树形结构的ASCII可视化功能，用于调试
"""

from typing import List, Optional, TYPE_CHECKING
from src.config import get_config

if TYPE_CHECKING:
    from src.index_tree import TreeNode, IndexTree


class TreeVisualizer:
    """树形结构可视化器"""
    
    def __init__(self):
        self.config = get_config().logging
    
    def format_tokens(self, tokens: List[int]) -> str:
        """
        格式化token列表显示
        :param tokens: token列表
        :return: 格式化后的字符串
        """
        if not tokens:
            return "[]"
        
        display_length = self.config.tree_debug_token_display_length
        
        if display_length == -1:
            # 显示全部
            return str(tokens)
        elif len(tokens) <= display_length:
            # 不需要截断
            return str(tokens)
        else:
            # 需要截断并添加省略号
            truncated = tokens[:display_length]
            return f"{truncated[:-1] + [truncated[-1]]}..."
    
    def format_node_info(self, node: 'TreeNode') -> str:
        """
        格式化节点信息
        :param node: 树节点
        :return: 格式化后的节点信息字符串
        """
        parts = []
        
        if self.config.tree_debug_show_hashes:
            parts.append(f"[{node.data['curr_hash']}]")
        
        token_str = self.format_tokens(node.data['token'])
        parts.append(f"tokens:{token_str}")
        
        return " ".join(parts)
    
    def generate_tree_ascii(self, tree: 'IndexTree', highlight_hash: Optional[int] = None) -> str:
        """
        生成树的ASCII可视化
        :param tree: IndexTree实例
        :param highlight_hash: 需要高亮的节点哈希（用于标记新增或修改的节点）
        :return: ASCII树形图字符串
        """
        if not tree.root:
            return "   (empty tree)"
        
        max_depth = self.config.tree_debug_max_depth
        max_nodes = self.config.tree_debug_max_nodes
        
        lines = []
        node_count = [0]  # 使用列表以便在递归中修改
        
        def traverse(node: 'TreeNode', prefix: str = "", is_last: bool = True, depth: int = 0):
            if depth >= max_depth or node_count[0] >= max_nodes:
                if depth >= max_depth:
                    lines.append(f"{prefix}└─ ... (max depth {max_depth} reached)")
                else:
                    lines.append(f"{prefix}└─ ... (max nodes {max_nodes} reached)")
                return
            
            node_count[0] += 1
            
            # 构建当前节点的显示行
            connector = "└─ " if is_last else "├─ "
            node_info = self.format_node_info(node)
            
            # 添加高亮标记
            highlight_mark = ""
            if highlight_hash and node.data['curr_hash'] == highlight_hash:
                highlight_mark = " ← NEW"
            
            lines.append(f"{prefix}{connector}{node_info}{highlight_mark}")
            
            # 为子节点准备前缀
            if node.child:
                child_prefix = prefix + ("    " if is_last else "│   ")
                for i, child in enumerate(node.child):
                    is_last_child = (i == len(node.child) - 1)
                    traverse(child, child_prefix, is_last_child, depth + 1)
        
        traverse(tree.root)
        return "\n".join(lines)
    
    @staticmethod
    def get_tree_stats(tree: 'IndexTree') -> dict:
        """
        获取树的统计信息
        :param tree: IndexTree实例
        :return: 包含统计信息的字典
        """
        if not tree.root:
            return {"total_nodes": 0, "max_depth": 0, "leaf_count": 0}
        
        stats = {"total_nodes": 0, "max_depth": 0, "leaf_count": 0}
        
        def traverse(node: 'TreeNode', depth: int = 0):
            stats["total_nodes"] += 1
            stats["max_depth"] = max(stats["max_depth"], depth)
            
            if not node.child:  # 叶子节点
                stats["leaf_count"] += 1
            else:
                for child in node.child:
                    traverse(child, depth + 1)
        
        traverse(tree.root)
        return stats
    
    def format_operation_info(self, operation: str, **kwargs) -> str:
        """
        格式化操作信息
        :param operation: 操作类型 (INSERT, DELETE, CLEAR)
        :param kwargs: 操作相关参数
        :return: 格式化后的操作信息
        """
        if operation == "INSERT":
            curr_hash = kwargs.get('curr_hash')
            pre_hash = kwargs.get('pre_hash')
            tokens = kwargs.get('tokens', [])
            token_str = self.format_tokens(tokens)
            return f"INSERT(curr_hash={curr_hash}, pre_hash={pre_hash}, tokens={token_str})"
        
        elif operation == "DELETE":
            curr_hash = kwargs.get('curr_hash')
            pre_hash = kwargs.get('pre_hash')
            return f"DELETE(curr_hash={curr_hash}, pre_hash={pre_hash})"
        
        elif operation == "CLEAR":
            return "CLEAR(all nodes removed)"
        
        else:
            return f"{operation}({kwargs})"
    
    def generate_tree_debug_log(
        self, tree: 'IndexTree', operation: str, highlight_hash: Optional[int] = None, **op_kwargs
    ) -> str:
        """
        生成完整的树调试日志
        :param tree: IndexTree实例
        :param operation: 操作类型
        :param highlight_hash: 需要高亮的节点哈希
        :param op_kwargs: 操作相关参数
        :return: 完整的调试日志字符串
        """
        if not self.config.tree_debug_enabled:
            return ""
        
        lines = []
        lines.append(f"Tree state after {operation} operation:")
        
        # 操作信息
        op_info = self.format_operation_info(operation, **op_kwargs)
        lines.append(f"├─ Operation: {op_info}")
        
        # 树统计信息
        stats = self.get_tree_stats(tree)
        lines.append(
            f"├─ Tree stats: {stats['total_nodes']} nodes, depth={stats['max_depth']}, leaves={stats['leaf_count']}"
        )
        
        # 树结构
        lines.append("└─ Structure:")
        tree_ascii = self.generate_tree_ascii(tree, highlight_hash)
        # 为树结构添加缩进
        tree_lines = tree_ascii.split('\n')
        for line in tree_lines:
            lines.append(f"   {line}")
        
        return "\n".join(lines)